import os

from langgraph.prebuilt import create_react_agent
from langchain_core.tools import tool
from typing import Annotated
from langgraph.graph import MessagesState
from langchain_community.utilities import SearchApiAPIWrapper
from server.common import chat_model
import requests
import json
from server.adaptive_rag.adaptive_rag import AdaptiveRag
from server.reminder_maker.reminder_maker import ReminderMaker

# from PIL import Image
# import io


# 假设img_data是bytes类型的图像数据
# img_data = adaptive_rag.get_graph().draw_mermaid_png()
# img_file = io.BytesIO(img_data)
# img = Image.open(img_file)
# img.show()


@tool
def query_kb(question: Annotated[str, "user's question"]):
    """This tool is used to search vectorstore. The vectorstore contains documents related to eye health.
    Use the vectorstore for questions on these topics."""
    adaptive_rag = AdaptiveRag().create_workflow()
    output = adaptive_rag.invoke({"question": question}, stream_mode="values")
    return output["generation"]


@tool
def create_reminder(task: Annotated[str, "Transformed sentence to a reminder task from user's request"]):
    """This tool is used to create a reminder and record it into database for user. Then return the result of the
    operation. If user wants you record something, or remind her/him something in the future, you can use this tool.
    """
    reminder_maker = ReminderMaker().create_workflow()
    output = reminder_maker.invoke({"task": task})
    return output["res"]


@tool
def web_search(query: Annotated[str, "The input of web search"]):
    """This tool is used to search the web page. Input the query and return the result of the search."""
    search = SearchApiAPIWrapper()
    print("query:" + query)
    res = search.run(query)
    print("results:" + str(res))
    return res


@tool
def weather_search(city: Annotated[str, "the name of city"]):
    """This tool is used to search the current weather of specified city. It will return the weather info of the
    city."""
    url = f"https://api.seniverse.com/v3/weather/now.json?key={os.getenv('SENIVERSE_API_KEY')}&location={city}&language=zh-Hans&unit=c"
    try:
        response = requests.get(url)
        if response.status_code != 200:
            return "request error, please try again."
        dict_res = json.loads(response.text)["results"][0]["now"]
        res = f"{city}的天气情况如下：天气：{dict_res['text']}，温度：{dict_res['temperature']}°C。"
        return res
    except Exception as e:
        print(f"请求出现错误: {e}")


class Agent:

    def __init__(self):
        self.agent = create_react_agent(
            chat_model.qwen, tools=[query_kb, web_search, create_reminder, weather_search]
        )

    def __call__(self, state: MessagesState):
        result = self.agent.invoke({"messages": state["messages"]})
        return result
